package com.lagou.edu.filter;

import com.lagou.edu.config.RateLimitConfig;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.couchbase.CouchbaseProperties;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 基于 Bucket4j 实现的限流过滤器 （令牌桶算法）
 */
@Slf4j
@Component  // 让容器扫描到，等同于注册了
@RefreshScope
public class RateLimitGatewayFilter implements GlobalFilter, Ordered {

    @Value("${lagou.ratelimit.capacity}")
    private int capacitySize;
    @Value("${lagou.ratelimit.refillTokens}")
    private int refillTokenSize;
    @Value("${lagou.ratelimit.refillDuration}")
    private int refillDurationSize;


//    int capacity = 1; //桶的最大容量，即能装载 Token 的最大数量
//    int refillTokens = 1; //每次 Token 补充量
//    Duration refillDuration = Duration.ofSeconds(1); //补充 Token 的时间间隔

    private static final Map<String, String> LAST_CACHE = new ConcurrentHashMap<>(); //上一次配置的缓存 IP
    private static final Map<String, Bucket> CACHE = new ConcurrentHashMap<>();

    /**
     * 这里指定了 bucket 的容量为 10 且每一秒会补充 2 个 Token
     * @return
     */
    private Bucket createNewBucket(int capacity, int refillTokens, int refillDurationSeconds) {
        Duration refillDuration = Duration.ofSeconds(refillDurationSeconds);
        Refill refill = Refill.of(refillTokens, refillDuration);
        Bandwidth limit = Bandwidth.classic(capacity, refill);
        return Bucket4j.builder().addLimit(limit).build();
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        // 从request对象中获取客户端ip
        String ip = exchange.getRequest().getRemoteAddress().getHostString();
        final String currentData = ip +"-" + capacitySize + "-" + refillTokenSize + "-" + refillDurationSize;
        log.info("config==> capacity:{}, refillTokens:{}, refillDuration:{}", capacitySize, refillTokenSize, refillDurationSize);
        Bucket bucket = CACHE.computeIfAbsent(ip, k -> createNewBucket(capacitySize, refillTokenSize, refillDurationSize));
        if(!LAST_CACHE.containsKey(ip) || !currentData.equals(LAST_CACHE.get(ip))) { //配置变了，需要更新流控规则
            bucket = createNewBucket(capacitySize, refillTokenSize, refillDurationSize);
            CACHE.put(ip, bucket);
            LAST_CACHE.put(ip, currentData);
            log.info("流控规则配置更新，现在更新流控组件规则");
        }
        log.info("IP: " + ip + ", TokenBucket Available Tokens: " + bucket.getAvailableTokens());
        if (bucket.tryConsume(1)) {
            return chain.filter(exchange);
        } else {
            exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
            return exchange.getResponse().setComplete();
        }
    }

    @Override
    public int getOrder() {
        return -1000;
    }
}
